import cv2
import numpy as np
import pyzbar.pyzbar as pyzbar

BASE_QR = 0x01
CV_QR = 0x02
ZBAR_QR = 0x03

class QR_Rec():
    def __init__(self, qr_mode):
        self.qr_mode = qr_mode
        self.qr_img = None

        # 设置检测器
        self.qrcoder = cv2.QRCodeDetector()

    def recQR(self, img=None):
        qr_roi = (0, 0, 0, 0)
        qr_str = ""

        if not img is None:
            self.qr_img = img
            if self.qr_mode == BASE_QR:
                qr_roi, qr_str = self.baseQR()
            elif self.qr_mode == CV_QR:
                qr_roi, qr_str = self.cvQR()
            elif self.qr_mode == ZBAR_QR:
                qr_roi, qr_str = self.zbarQR()
        else:
            print("图像为空！！")

        return qr_roi, qr_str

    def baseQR(self):
        img_gray = cv2.cvtColor(self.qr_img, cv2.COLOR_BGR2GRAY)

        # 二值化
        ret, thr_img = cv2.threshold(img_gray, 200, 255, cv2.THRESH_BINARY_INV)
        cv2.imshow("thr_img", thr_img)

        # 膨胀
        kernel = np.ones((5, 5), np.uint8)
        dilate = cv2.dilate(thr_img, kernel, iterations=12)
        cv2.imshow("dilat", dilate)

        # 腐蚀
        kernel = np.ones((5, 5), np.uint8)
        erosion = cv2.erode(dilate, kernel, iterations=12)
        cv2.imshow("erode", erosion)

        # 查找轮廓
        contours, hierarchy = cv2.findContours(erosion, cv2.RETR_TREE, \
                                               cv2.CHAIN_APPROX_SIMPLE)[-2:]
        x, y, w, h = cv2.boundingRect(contours[0])
        qr_roi = self.qr_img[y:y + h, x:x + w]
        self.qr_img = qr_roi
        return (x, y, w, h), self.cvQR()[1]


    def cvQR(self):
        qr_rect = (0, 0, 0, 0)
        codeinfo = ""

        try:
            # 检测识别二维码
            codeinfo, points, straight_qrcode = self.qrcoder.detectAndDecode(self.qr_img)
            qr_rect = cv2.boundingRect(points)
        except:
            pass

        return qr_rect, codeinfo


    def zbarQR(self):
        qr_rect = (0, 0, 0, 0)
        qr_str = ""
        barcodes = pyzbar.decode(self.qr_img)

        for barcode in barcodes:
            # 解析二维码内容
            qr_str = barcode.data.decode("utf-8")

            # 解析二维码所处图片中的像素位置
            qr_rect = barcode.rect

        return qr_rect, qr_str